#include "solver.h"
#include <fstream>


#define block(i, j, lda) A->block[(j) * (lda) + (i)]
#define tile(i, j, lda) A->tile[(j) * (lda) + (i)]
#define a(i, j, lda) A->a[(j) * (lda) + (i)]



#define BUFMAX 256

void destroy_T(TILE *T)
{
    if (!T)
        return;
    my_free(T->a, T->m * T->n, sizeof(double));
    my_free(T, 1, sizeof(TILE));
}

void destroy_B(BLOCK *B)
{
    if (!B)
        return;
    for (int i = 0; i < B->m; i++)
    {
        for (int j = 0; j < B->n; j++)
        {
            destroy_T(B->tile[i * B->n + j]);
        }
    }
    my_free(B->tile, B->m * B->n, sizeof(TILE *));
    my_free(B, 1, sizeof(BLOCK));
}

void destroy_M(MATRIX *A)
{
    if (!A)
        return;
    for (int i = 0; i < A->m; i++)
    {
        for (int j = 0; j < A->n; j++)
        {
            destroy_B(A->block[i * A->n + j]);
        }
    }
    if (A->name)
        my_free(A->name, strlen(A->name) + 1, sizeof(char));
    my_free(A->a, A->real_nonz, sizeof(double));
    my_free(A->row, A->real_m + 1, sizeof(int));
    my_free(A->col, A->real_nonz, sizeof(int));
    my_free(A, 1, sizeof(MATRIX));
}

TILE *construct_T(int m, int n)
{
    TILE *t = (TILE *)my_malloc(1, sizeof(TILE));

    t->m = m;
    t->n = n;
    t->nonz = 0;

    t->a = (double *)my_malloc(m * n, sizeof(double));
    return t;
}

BLOCK *construct_B(int m, int n)
{
    BLOCK *b = (BLOCK *)my_malloc(1, sizeof(BLOCK));

    b->m = m;
    b->n = n;
    b->nonz = 0;

    b->tile = (TILE **)my_malloc(m * n, sizeof(TILE *));
    return b;
}

MATRIX *construct_M(void)
{
    MATRIX *A = (MATRIX *)my_malloc(1, sizeof(MATRIX));

    A->type = NONE;
    A->name = NULL;
    A->a = NULL;
    A->row = NULL;
    A->col = NULL;
    A->real_m = 0;
    A->real_n = 0;
    A->real_nonz = 0;

    A->block_size = 0;
    A->tile_size = 0;

    A->threshold = 0;

    A->m = 0;
    A->n = 0;
    A->nonz = 0;
    A->block = NULL;
    A->output_path = NULL;

    return A;
}


void formfull(MATRIX* A){
    if (!A->symmetric) return;

    int *tmp = (int *)malloc(A->real_n * sizeof(int));

    // 转置矩阵
    double *t_a = (double *)malloc(A->real_nonz * sizeof(double));
    int *t_row = (int *)malloc(A->real_nonz * sizeof(int));
    int *t_col = (int *)malloc((A->real_n + 1) * sizeof(int));

    // 标记每列元素数量
    memset(tmp, 0, A->real_n * sizeof(int));
    for (int j = 0; j < A->real_n; j++) {
        for (int i = A->col[j]; i < A->col[j + 1]; i++) {
            tmp[ A->row[i] ]++;
        }
    }
    // 设置转置矩阵列指针
    t_col[0] = 0;
    for (int j = 0; j < A->real_n; j++) {
        t_col[j + 1] = t_col[j] + tmp[j];
        tmp[j] = t_col[j];
    }
    // 矩阵转置
    for (int j = 0; j < A->real_n; j++) {
        for (int i = A->col[j]; i < A->col[j + 1]; i++) {
            int col_index = A->row[i];
            t_row[ tmp[col_index] ] = j;
            t_a  [ tmp[col_index] ] = A->a[i];
            tmp[col_index]++;
        }
    }

    // 设置新矩阵
    int new_nonz = A->real_nonz * 2 - A->real_n;
    double *a2 = (double *)malloc(new_nonz * sizeof(double));
    int *row2 = (int *)malloc(new_nonz * sizeof(int));
    int *col2 = (int *)malloc((A->real_n + 1) * sizeof(int));

    // 矩阵相加
    col2[0] = 0;
    int k = 0;
    for (int j = 0; j < A->real_n; j++) {
        // 上三角部分
        for (int i = t_col[j]; i < t_col[j + 1]; i++) {
            if (t_row[i] != j) {
                row2[k] = t_row[i];
                a2[k] = t_a[i];
                k++;
            }
        }
        // 下三角部分
        for (int i = A->col[j]; i < A->col[j + 1]; i++) {
            row2[k] = A->row[i];
            a2[k] = A->a[i];
            k++;
        }
        col2[j + 1] = k;
    }

    free(tmp);
    free(t_a);
    free(t_row);
    free(t_col);
    free(A->a);
    free(A->row);
    free(A->col);

    A->a = a2;
    A->row = row2;
    A->col = col2;
    A->real_nonz = new_nonz;
}

// 读取为CSC格式保存
void readrb(const char *file, MATRIX *A){
    FILE *fp = fopen(file, "r");
    if (!fp) ABORT("fopen fail");

    A->type = CSC;
    char buf[BUFMAX], *p;
    int tmp;

     // line1
     p = fgets(buf, BUFMAX, fp);

     // line2
     p = fgets(buf, BUFMAX, fp);
     int lines[4];
     sscanf(buf, "%d %d %d %d", &lines[0], &lines[1], &lines[2], &lines[3]);
 
     // line 3
     p = fgets(buf, BUFMAX, fp);
     char type[3];
     sscanf(buf, "%3c %d %d %d %d", type, &A->real_m, &A->real_n, &A->real_nonz, &tmp);
     if (A->real_m != A->real_n) ABORT("not square matrix");

     if(type[0] != 'r' && type[0] != 'R') {
        ABORT("The matrix is not real type");
    }
    if(type[2] != 'a' && type[2] != 'A') {
        ABORT("The matrxi is not an assembled matrix");
    }
    int symmetric = type[1] == 's' || type[1] == 'S';
    A->symmetric = symmetric;

    p = fgets(buf, BUFMAX, fp);
    
    int n = A->real_n;
    int nonz = A->real_nonz;
    
    A->a = (double*)malloc(nonz * sizeof(double));
    A->row = (int*)malloc(nonz * sizeof(int));
    A->col = (int*)malloc((n+1) * sizeof(int));

    // read col
    for (int i = 0; i <= n; i++) {
        tmp = fscanf(fp, "%d", A->col + i);
        A->col[i]--; // base-0
    }

    // read row
    for (int i = 0; i < nonz; i++) {
        tmp = fscanf(fp, "%d", A->row + i);
        A->row[i]--;  // base-0
    }

    // read val
    for (int i = 0; i < nonz; i++) {
        tmp = fscanf(fp, "%lf", A->a + i);
    }
    fclose(fp);

    formfull(A);

    // 保存矩阵名称
    int s, e;
    for (int i = strlen(file) - 1; i >= 0; i--)
    {
        if (file[i] == '/')
        {
            s = i;
            break;
        }
        if (file[i] == '.')
        {
            e = i;
        }
    }
    // TODO: BUG e < s
    A->name = (char *)my_malloc(e - s, sizeof(char));
    for (int i = 0; i < e - s - 1; i++)
    {
        A->name[i] = file[s + 1 + i];
    }
    A->name[e - s - 1] = '\0';

}

void readcsc(const char *file, MATRIX *A){
    std::ifstream fin;
    fin.open(file);
    if(!fin.good()){
        printf("Open file Failed. at %s,%d\n", __FILE__ ,__LINE__);
    }

    int colptr_len;
    int rowind_len;
    int data_len;
    fin.read((char*)&A->real_n, sizeof(int));
    A->real_m = A->real_n;
    fin.read((char*)&A->real_nonz, sizeof(int));
    

    printf("m = %d,  nonz = %d\n", A->real_n, A->real_nonz);


    // malloc
    A->row = (int *)malloc(A->real_nonz * sizeof(int));
    A->col = (int *)malloc((A->real_n + 1) * sizeof(int));
    A->a = (double *)malloc(A->real_nonz * sizeof(double));


    fin.read((char*)&colptr_len, sizeof(int));
    fin.read((char*)A->col, sizeof(int)*colptr_len);


    fin.read((char*)&rowind_len, sizeof(int));
    fin.read((char*)A->row, sizeof(int)*rowind_len);

  
    fin.read((char*)&data_len, sizeof(int));
    printf("data_len = %d\n", data_len);
    fin.read((char*)A->a, sizeof(double)*data_len);


    // col 每个元素 -1
    for(int i = 0; i < A->real_n+1; i++){
        A->col[i]--;
    }
    for(int i = 0; i < A->real_nonz; i++){
        A->row[i]--;
    }

    A->type = CSC;

     // 保存矩阵名称
     int s, e;
     for (int i = strlen(file) - 1; i >= 0; i--)
     {
         if (file[i] == '/')
         {
             s = i;
             break;
         }
         if (file[i] == '.')
         {
             e = i;
         }
     }
     // TODO: BUG e < s
     A->name = (char *)my_malloc(e - s, sizeof(char));
     for (int i = 0; i < e - s - 1; i++)
     {
         A->name[i] = file[s + 1 + i];
     }
     A->name[e - s - 1] = '\0';
}


void read_to_csr(const char *filename, MATRIX *A)
{
    if (!filename || !A)
        ABORT("NULL pointer");

    dreadrb(filename, &A->real_m, &A->real_n, &A->real_nonz, &A->a, &A->col, &A->row);
    if (A->real_m != A->real_n)
        ABORT("A is not square");

    A->type = CSR;

    int s, e;
    for (int i = strlen(filename) - 1; i >= 0; i--)
    {
        if (filename[i] == '/')
        {
            s = i;
            break;
        }
        if (filename[i] == '.')
        {
            e = i;
        }
    }
    // TODO: BUG e < s
    A->name = (char *)my_malloc(e - s, sizeof(char));
    for (int i = 0; i < e - s - 1; i++)
    {
        A->name[i] = filename[s + 1 + i];
    }
    A->name[e - s - 1] = '\0';
}

void init(MATRIX *A, int block_size, int tile_size)
{
    if (!A)
        return;
    if (A->type == CSR){
        printf("A format is CSR\n");
    

        set_size(A, block_size, tile_size);

        printf("block_size = %d, tile_size = %d\n", A->block_size, A->tile_size);

        A->m = A->real_m / (A->block_size * A->tile_size);
        if (A->real_m % (A->block_size * A->tile_size) != 0)
        {
            A->m++;
        }
        A->n = A->m;
        A->nonz = 0;
        A->block = (BLOCK **)my_malloc(A->m * A->n, sizeof(BLOCK *));

        int lda = A->n;

        int bi, bj;
        int ti, tj;
        int ri, rj;

        // convert A from CSR to BT
        for (int i = 0; i < A->real_m; i++)
        {
            get_position(A, i, &bi, &ti, &ri);
            for (int j = A->row[i]; j < A->row[i + 1]; j++)
            {
                if (A->a[j] == 0)
                    continue; // rb file may have zero element
                get_position(A, A->col[j], &bj, &tj, &rj);

                int bm = get_bs(A, bi);
                int bn = get_bs(A, bj);
                int ldb = bn;
                if (!A->block[bi * lda + bj])
                {
                    A->block[bi * lda + bj] = construct_B(bm, bn);
                    A->nonz++;
                    A->num_block++;
                }
                int tm = get_ts(A, bi, ti);
                int tn = get_ts(A, bj, tj);
                int ldt = tn;
                if (!A->block[bi * lda + bj]->tile[ti * ldb + tj])
                {
                    A->block[bi * lda + bj]->tile[ti * ldb + tj] = construct_T(tm, tn);
                    A->block[bi * lda + bj]->nonz++;
                    A->num_tile++;
                }
                A->block[bi * lda + bj]->tile[ti * ldb + tj]->a[ri * ldt + rj] = A->a[j];
                A->block[bi * lda + bj]->tile[ti * ldb + tj]->nonz++;
            }
        }

        A->type = DENSE;

        A->mem = get_memory() / (1024 * 1024);
    }

    // CSC稀疏格式转换为BLOCK_TILE格式
    else if(A->type == CSC){
        set_size(A, block_size, tile_size);
        printf("A format is CSC\n");
        printf("block_size = %d, tile_size = %d\n", A->block_size, A->tile_size);

        A->m = A->real_m / (A->block_size * A->tile_size);
        if (A->real_m % (A->block_size * A->tile_size) != 0)
        {
            A->m++;
        }
        A->n = A->m;
        A->nonz = 0;
        A->block = (BLOCK **)my_malloc(A->m * A->n, sizeof(BLOCK *));

        int lda = A->n;

        int bi, bj;
        int ti, tj;
        int ri, rj;

        // convert A from CSC to BT
        for (int j = 0; j < A->real_m; j++)
        {
            get_position(A, j, &bj, &tj, &rj);
            for (int i = A->col[j]; i < A->col[j + 1]; i++)
            {
                if (A->a[i] == 0)
                    continue; // rb file may have zero element
                get_position(A, A->row[i], &bi, &ti, &ri);

                int bm = get_bs(A, bj);
                int bn = get_bs(A, bi);
                int ldb = bn;
                if (!A->block[bj * lda + bi])
                {
                    A->block[bj * lda + bi] = construct_B(bm, bn);
                    A->nonz++;
                    A->num_block++;
                }
                int tm = get_ts(A, bj, tj);
                int tn = get_ts(A, bi, ti);
                int ldt = tn;
                if (!A->block[bj * lda + bi]->tile[tj * ldb + ti])
                {
                    A->block[bj * lda + bi]->tile[tj * ldb + ti] = construct_T(tm, tn);
                    A->block[bj * lda + bi]->nonz++;
                    A->num_tile++;
                }
                A->block[bj * lda + bi]->tile[tj * ldb + ti]->a[rj * ldt + ri] = A->a[i];
                A->block[bj * lda + bi]->tile[tj * ldb + ti]->nonz++;
            }
        }

        A->type = DENSE;

        A->mem = get_memory() / (1024 * 1024);
    }
}
